from ultralytics import YOLO

# 立新训练样本：训练正面缺角检测模型
# 图片尺寸：1280 * 1024px
if __name__ == '__main__':
    model = YOLO("yolo11n.yaml").load("yolo11n.pt")  # build from YAML and transfer weights
    # model = YOLO("runs/detect/train/weights/best.pt")

    # Train the model
    results = model.train(data="datasets/lixin_front_chipping/data.yaml",  # 数据集配置文件
                          epochs=1,                         # 总共训练的轮数
                          imgsz=1280,                        # 输入图像尺寸（训练和验证）
                          batch=3,                           # 每个批次的图像数量
                          device=0,                          # 使用第0号GPU进行训练
                          workers=0,                         # 数据加载的线程数
                          close_mosaic=10,                   # 最后10轮关闭 Mosaic 数据增强
                          single_cls=True,                  # 所有目标不视为同一类别
                          rect=True,                         # 使用矩形推理，保留原始宽高比
                          warmup_epochs=2,                   # 前2轮进行学习率预热
                          augment=True,                      # 启用数据增强
                          )

    # Evaluate the model's performance on the validation set
    results = model.val(device=0,
                        imgsz=1280,
                        rect=True,
                        conf=0.1,
                        iou=0.5,
                        max_det=300,
                        plots=True)

    # Export the model to ONNX format
    success = model.export(format="onnx", imgsz=[1024, 1280])
